import zmq
import numpy as np
import time
import sys
import traceback
import cv2
import os
from datetime import datetime
from typing import Dict, Any

def echo_exception():
    """获取并格式化当前异常的详细信息"""
    exc_type, exc_value, exc_traceback = sys.exc_info()
    tb_lines = traceback.format_exception(exc_type, exc_value, exc_traceback)
    return "".join(tb_lines)

class TimestampRecorder:
    """记录时间戳的工具类"""
    def __init__(self):
        self.start_time = time.time()
        print(f"服务器启动于 {self.start_time} 秒")

    def get_current_timestamp(self):
        current_time = time.time() - self.start_time
        return round(current_time, 6)

class ImageZmqServer:
    """图像ZMQ服务器,用于接收和处理图像数据"""
    def __init__(
        self,
        zmq_ip: str,
        zmq_port: int,
        no_cmd_timeout: float = -1.0,  # -1 表示无超时限制
        save_dir: str = "doc/my_data"
    ):
        self.recorder = TimestampRecorder()
        
        # 初始化ZMQ通信
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.REP)
        self.socket.bind(f"tcp://{zmq_ip}:{zmq_port}")
        
        self.poller = zmq.Poller()
        self.poller.register(self.socket, zmq.POLLIN)
        
        self.zmq_ip = zmq_ip
        self.zmq_port = zmq_port
        self.no_cmd_timeout = no_cmd_timeout
        
        # 存储最新接收的图像
        self.latest_image = None
        self.latest_timestamp = None
        
        # 存储最新接收的文本
        self.latest_text = None
        self.latest_text_language = "auto"
        self.latest_text_timestamp = None
        
        # 存储最新的坐标数据
        self.coords_data = None
        self.coords_timestamp = None
        
        # 图像保存目录
        self.save_dir = save_dir
        # 确保保存目录存在
        os.makedirs(self.save_dir, exist_ok=True)
        
        print(f"图像ZMQ服务器启动在 {self.zmq_ip}:{self.zmq_port}")

    def run(self):
        """运行服务器主循环"""
        while True:
            try:
                # 等待客户端请求，超时时间为no_cmd_timeout
                if self.no_cmd_timeout <= 0:
                    # 无超时限制
                    socks = dict(self.poller.poll())
                else:
                    # 有超时限制
                    socks = dict(self.poller.poll(int(self.no_cmd_timeout * 1000)))
                    
                if self.socket in socks and socks[self.socket] == zmq.POLLIN:
                    msg = self.socket.recv_pyobj()
                else:
                    # 超时继续循环
                    continue
            except KeyboardInterrupt:
                # 用户中断
                break
            except Exception as e:
                exception_str = echo_exception()
                print(f"错误: {exception_str}")
                continue

            try:
                # 检查消息格式
                if not isinstance(msg, dict):
                    print(f"错误: 收到无效消息，已忽略")
                    self.socket.send_pyobj(
                        {
                            "cmd": "UNKNOWN",
                            "data": "错误: 收到无效消息",
                        }
                    )
                    continue

                # 处理不同类型的命令
                if msg["cmd"] == "PING":
                    # 心跳检测
                    print(f"收到PING消息")
                    reply_msg = {
                        "cmd": "PING",
                        "data": {
                            "timestamp": self.recorder.get_current_timestamp(),
                            "status": "OK",
                        },
                    }
                    self.socket.send_pyobj(reply_msg)
                
                elif msg["cmd"] == "SEND_IMAGE":
                    # 接收图像数据
                    print(f"收到SEND_IMAGE消息")
                    image_data = msg["data"]["image"]
                    self.latest_image = image_data
                    self.latest_timestamp = self.recorder.get_current_timestamp()
                    
                    # # 显示接收到的图像
                    # if self.latest_image is not None:
                    #     cv2.imshow("Image", self.latest_image)
                    #     cv2.waitKey(1)  # 非阻塞更新显示窗口
                        
                    # 保存图像到指定目录
                    if isinstance(self.latest_image, np.ndarray):
                        filename = f"{self.save_dir}/receive_image.jpg"
                        try:
                            cv2.imwrite(filename, self.latest_image)
                            print(f"图像已保存到: {filename}")
                        except Exception as e:
                            print(f"保存图像失败: {e}")
                    
                    # 回复确认消息
                    reply_msg = {
                        "cmd": "SEND_IMAGE",
                        "data": {
                            "timestamp": self.latest_timestamp,
                            "status": "OK",
                            "image_size": image_data.shape if hasattr(image_data, 'shape') else None
                        },
                    }
                    self.socket.send_pyobj(reply_msg)
                
                elif msg["cmd"] == "SEND_TEXT":
                    # 接收文本数据
                    print(f"收到SEND_TEXT消息")
                    text_data = msg["data"]["text"]
                    language = msg["data"].get("language", "auto")
                    self.latest_text = text_data
                    self.latest_text_language = language
                    self.latest_text_timestamp = self.recorder.get_current_timestamp()
                    
                    # 打印接收到的文本
                    print(f"接收到文本: {text_data} (语言: {language})")
                    
                    # 回复确认消息
                    reply_msg = {
                        "cmd": "SEND_TEXT",
                        "data": {
                            "timestamp": self.latest_text_timestamp,
                            "status": "OK",
                            "text_length": len(text_data) if text_data else 0,
                            "language": language
                        },
                    }
                    self.socket.send_pyobj(reply_msg)
                
                elif msg["cmd"] == "GET_TEXT":
                    # 返回最新的文本
                    print(f"收到GET_TEXT消息")
                    reply_msg = {
                        "cmd": "GET_TEXT",
                        "data": {
                            "timestamp": self.latest_text_timestamp,
                            "text": self.latest_text,
                            "language": self.latest_text_language
                        },
                    }
                    self.socket.send_pyobj(reply_msg)
                
                elif msg["cmd"] == "GET_COORDS":
                    # 返回最新的坐标数据
                    if self.coords_data is not None:
                        # 如果有坐标数据，返回完整的数据结构
                        reply_msg = {
                            "cmd": "GET_COORDS",
                            "data": {
                                "bbox_list": self.coords_data.get("bbox_list", []),
                                "found_bbox": self.coords_data.get("found_bbox", False),
                                "timestamp": self.coords_timestamp,
                            },
                        }
                        # 发送后清除坐标数据，防止下次请求返回相同的坐标
                        temp_coords = self.coords_data
                        self.coords_data = None
                        self.coords_timestamp = None
                        
                        self.socket.send_pyobj(reply_msg)
                        
                        # 打印返回的坐标信息
                        print(f"已返回坐标 {temp_coords} 并清除缓存")
                    else:
                        # 如果没有坐标数据，返回空数据
                        reply_msg = {
                            "cmd": "GET_COORDS",
                            "data": {
                                "bbox_list": [],
                                "found_bbox": False,
                                "timestamp": None,
                            },
                        }
                        self.socket.send_pyobj(reply_msg)
                
                elif msg["cmd"] == "GET_IMAGE":
                    # 返回最新的图像
                    print(f"收到GET_IMAGE消息")
                    reply_msg = {
                        "cmd": "GET_IMAGE",
                        "data": {
                            "timestamp": self.latest_timestamp,
                            "image": self.latest_image,
                        },
                    }
                    self.socket.send_pyobj(reply_msg)
                
                else:
                    # 未知命令
                    print(f"收到未知命令: {msg['cmd']}")
                    self.socket.send_pyobj(
                        {
                            "cmd": msg["cmd"],
                            "data": f"错误: 未知命令 {msg['cmd']}",
                        }
                    )

            except Exception as e:
                exception_str = echo_exception()
                try:
                    self.socket.send_pyobj(
                        {"cmd": msg["cmd"], "data": f"错误: {exception_str}"}
                    )
                except:
                    pass
                print(f"错误: {exception_str}")
                continue

    def send_coords(self, bbox_list, found_bbox=True):
        """设置bbox坐标数据并更新时间戳
        
        Args:
            bbox_list: bbox坐标列表,每个元素为[x1, y1, x2, y2, label]
            found_bbox: 是否找到bbox的标志位
            
        Returns:
            bool: 设置成功返回True,否则返回False
        """
        # 设置服务器的坐标数据属性
        self.coords_data = {
            "bbox_list": bbox_list,
            "found_bbox": found_bbox
        }
        self.coords_timestamp = self.recorder.get_current_timestamp()
        
        if found_bbox and bbox_list:
            print(f"已设置bbox坐标列表: {bbox_list}, 标志位: {found_bbox}")
        else:
            print(f"未找到bbox坐标,发送默认值,标志位: {found_bbox}")
        return True

    def __del__(self):
        """清理资源"""
        try:
            self.socket.close()
            self.context.term()
            cv2.destroyAllWindows()  # 关闭所有OpenCV窗口
        except:
            pass

# 如果直接运行此脚本，启动服务器
if __name__ == "__main__":
    # 不需要从当前文件导入自身
    server = ImageZmqServer('127.0.0.1', 5555, save_dir='doc/my_data')
    server.run()